"""
现有网络模型的使用与修改
"""
import torchvision
from torch import nn

train_data = torchvision.datasets.CIFAR10("../data", train=True, download=True,
                                          transform=torchvision.transforms.ToTensor())

# 初始参数
# vgg16_false = torchvision.models.vgg16(pretrained=False)
# 训练好的参数
vgg16_true = torchvision.models.vgg16(pretrained=True)
print(vgg16_true)

# 在vgg16的classifier的基础上加上一个线性结构
vgg16_true.classifier.add_module('add_linear', nn.Linear(1000, 10))
print(vgg16_true)

# 修改vgg16的classifier的第6个层
vgg16_true.classifier[6] = nn.Linear(4096, 10)
print(vgg16_true)